import argparse
from ultralytics import YOLO
def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--imgsz', type=int, default=640, help='Image size')
    parser.add_argument('--batch', type=int, default=16, help='Batch size')
    parser.add_argument('--epochs', type=int, default=300, help='Number of epochs')
    parser.add_argument('--data', type=str, default='data.yaml', help='Path to dataset.yaml')
    parser.add_argument('--weights', type=str, default='YOLOv10x.pt', help='Initial weights path')
    parser.add_argument('--name', type=str, default='yolov10_bug', help='Name for the run')
    return parser.parse_args()

def main(opt):
    model = YOLO(opt.weights)
    model.train(
        data=opt.data,
        epochs=opt.epochs,
        batch=opt.batch,
        imgsz=opt.imgsz,  # Corrected from img to imgsz
        device='cuda',  # 设置为 'cuda' 使用 GPU
        name=opt.name
    )

if __name__ == "__main__":
    opt = parse_opt()
    main(opt)
